from app_config import get_pro
import pandas as pd
from app_config import get_engine_ts
import numpy as np


def index_weight():
    trade_date = "20240930"
    three_year_start = '20210801'
    three_year_end = '20240731'
    pre_000300_index = get_pro().index_weight(index_code='000015.SH', start_date=trade_date,
                                              end_date=trade_date)

    hs300_ts_code = pre_000300_index['con_code'].tolist()

    query_dividend = f"""
          SELECT ts_code, pay_date, cash_div_tax, base_share FROM dividend WHERE div_proc = '实施' AND cash_div_tax > 0 
          AND pay_date >= '{three_year_start}' AND pay_date <= '{three_year_end}'
          """
    dividend_data = pd.read_sql_query(query_dividend, get_engine_ts())
    dividend_data['cash_amount'] = dividend_data['cash_div_tax'] * dividend_data['base_share']

    grouped_df_total_cash = dividend_data.groupby('ts_code')['cash_amount'].sum().reset_index()
    grouped_df_total_cash.columns = ['ts_code', 'cash_amount_total']

    pre_000300_index = pd.merge(pre_000300_index, grouped_df_total_cash, left_on='con_code', right_on='ts_code',
                                how='left')

    # 构建查询语句
    query = f"""
       SELECT ts_code,total_mv FROM `daily_basic`
       WHERE ts_code IN ({','.join(f"'{code}'" for code in hs300_ts_code)})
       AND trade_date = '{three_year_end}'
       """
    ts_codes_total_mv = pd.read_sql_query(query, get_engine_ts())
    pre_000300_index = pd.merge(pre_000300_index, ts_codes_total_mv, on='ts_code', how='left')

    pre_000300_index["total_mv亿元"] = pre_000300_index['total_mv'] / 10000

    pre_000300_index['股息率'] = pre_000300_index['cash_amount_total'] / pre_000300_index['total_mv']

    amount_sum = pre_000300_index['股息率'].sum()
    pre_000300_index['amount_precent'] = pre_000300_index['股息率'] / amount_sum * 100

    # # 构建查询语句
    # query = f"""
    #        SELECT ts_code,close,total_share,free_share FROM `daily_basic`
    #        WHERE ts_code IN ({','.join(f"'{code}'" for code in hs300_ts_code)})
    #        AND trade_date = '{trade_date}'
    #        """
    # # 执行查询并将结果转换为DataFrame
    # stock_share = pd.read_sql_query(query, get_engine_ts())
    # hs300 = pd.merge(pre_000300_index, stock_share, left_on="con_code", right_on="ts_code", how='left')
    #
    # hs300['calc_percent'] = hs300['free_share'] / hs300['total_share'] * 100
    #
    # # 定义生成 b 列的规则
    # def generate_b(value):
    #     if value <= 15:
    #         return np.ceil(value)  # 上调至最接近的整数
    #     elif 15 < value <= 20:
    #         return 20
    #     elif 20 < value <= 30:
    #         return 30
    #     elif 30 < value <= 40:
    #         return 40
    #     elif 40 < value <= 50:
    #         return 50
    #     elif 50 < value <= 60:
    #         return 60
    #     elif 60 < value <= 70:
    #         return 70
    #     elif 70 < value <= 80:
    #         return 80
    #     else:
    #         return 100
    #
    # hs300['adjust_percent'] = hs300['calc_percent'].apply(generate_b)
    # hs300['adjust_share'] = hs300['adjust_percent'] / 100 * hs300['total_share']
    # hs300['adjust_amount'] = hs300['adjust_share'] * hs300['close']
    #
    # amount_sum = hs300['adjust_amount'].sum()
    # hs300['amount_precent'] = hs300['adjust_amount'] / amount_sum * 100

    pre_000300_index.to_excel("v_weight_000015上证红利指数.xlsx")


if __name__ == '__main__':
    index_weight()
